#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/ADT/GraphTraits.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/Analysis/IteratedDominanceFrontier.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/DebugInfoMetadata.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Use.h"
#include "llvm/IR/Value.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/ProfileData/InstrProf.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"
#include <cassert>
#include <cstddef>
#include <utility>
#include <vector>
#include <set>
#include <queue>
#include <unordered_map>

using namespace llvm;

#define DEBUG_TYPE "loop-search"
// STATISTIC(LoopSearched, "Number of loops has been found");
#define DEBUG 1


enum ColorType
{
    WHITE = 0,
    GRAY,
    BLACK
};

namespace llvm
{
    FunctionPass *createLSPass();
    void initializeLSPassPass(PassRegistry &);
} // namespace llvm

namespace
{
    struct LoopEntry
    {
        unsigned int Level;
        DomTreeNode *UpperLoop;                     //该循环的上层循环在支配树中的指针, 如无上层循环, 则为 nullptr
        std::vector<DomTreeNode *> Pred;            //从下层指向该循环入口的节点的指针
        std::vector<DomTreeNode *> SuccLoop;        //指向直接下层循环的指针
        std::string Tag;                            //循环 Tag
    };

    struct LSPass : public FunctionPass
    {
        static char ID; // Pass identification, replacement for typeid

        LSPass() : FunctionPass(ID)
        {
            initializeLSPassPass(*PassRegistry::getPassRegistry());
        }

        std::unordered_map<DomTreeNode *, ColorType> Visited;
        std::unordered_map<DomTreeNode *, LoopEntry> LE;            //循环入口
        std::queue<DomTreeNode *> Queue;                            //用于 BFS 的 Quene
        std::vector<DomTreeNode *> UpperMost_Loop;                  //最上层循环列表

        inline void queue_Clear()
        {
            while (!Queue.empty())
            {
                Queue.pop();
            }
        }

        //初始化 Visited, 使得 Visited 中含有 DT 中的所有节点
        void init_Map(DominatorTree &DT)
        {
            queue_Clear();
            Queue.push(DT.getRootNode());
            while (!Queue.empty())
            {
                DomTreeNode *Temp = Queue.front();
                Visited[Temp] = ColorType::WHITE;
                Queue.pop();
                BasicBlock *BB = Temp->getBlock();
                for (auto sSucc = succ_begin(BB), eSucc = succ_end(BB); sSucc != eSucc; sSucc++)
                {
                    DomTreeNode *Succ_Node = DT.getNode(*sSucc);
                    if (Visited.find(Succ_Node) != Visited.end())
                    {
                        Queue.push(Succ_Node);
                    }
                }
            }
        }

        //Reset Visited = ColorType::WHITE
        void visited_Clear()
        {
            for (auto Cur = Visited.begin(), End = Visited.end(); Cur != End; ++Cur)
            {
                Cur->second = ColorType::WHITE;
            }
        }

        DomTreeNode *get_Upper_Loop_Entry(DominatorTree &DT, DomTreeNode *LoopEntry)
        {
            queue_Clear();
            visited_Clear();
            Queue.push(LoopEntry);
            Visited[LoopEntry] = ColorType::GRAY;
            while (!Queue.empty())
            {
                DomTreeNode *CurNode = Queue.front();
                BasicBlock *Cur_BB = CurNode->getBlock();
                Queue.pop();
                for (auto sSucc = succ_begin(Cur_BB), eSucc = succ_end(Cur_BB); sSucc != eSucc; ++sSucc)
                {
                    DomTreeNode *Succ_Node = DT.getNode(*sSucc);
                    if (Visited[Succ_Node] != ColorType::WHITE)
                    {
                        continue;
                    }
                    else if (LE.find(Succ_Node) != LE.end() && DT.dominates(Succ_Node, LoopEntry))
                    {
                        return Succ_Node;
                    }
                    Visited[Succ_Node] = ColorType::GRAY;
                    Queue.push(Succ_Node);
                }
                Visited[CurNode] = ColorType::BLACK;
            }
            return nullptr;
        }

        static bool Compare(DomTreeNode *Node_1, DomTreeNode *Node_2)
        {
            return (Node_1->getLevel() < Node_2->getLevel());
        }

        //为每个循环的入口, 都找到其直接下层循环
        void get_Imm_Loop_Entry()
        {
            for (auto loopEntry = LE.begin(); loopEntry != LE.end(); ++loopEntry)
            {
                auto UpperLoop = loopEntry->second.UpperLoop;
                if (UpperLoop != nullptr)
                {
                    LE[UpperLoop].SuccLoop.push_back(loopEntry->first);
                }
                else
                {
                    UpperMost_Loop.push_back(loopEntry->first);
                }
            }
            //对每个 LoopEntry, 将其直接下层的循环按照 Level 进行排序, 得到数据流中循环依次执行的顺序
            for (auto loopEntry = LE.begin(); loopEntry != LE.end(); ++loopEntry)
            {
                std::sort(loopEntry->second.SuccLoop.begin(), loopEntry->second.SuccLoop.end(), Compare);
            }
            std::sort(UpperMost_Loop.begin(), UpperMost_Loop.end(), Compare);
        }

        void get_Loop_Tag(DomTreeNode *UpperLoopEntry, std::vector<DomTreeNode *> &SuccEntries)
        {
            int loop_num = 1;
            for (auto loopEntry = SuccEntries.begin(); loopEntry != SuccEntries.end(); ++loopEntry, ++loop_num)
            {
                if (UpperLoopEntry == nullptr)
                {
                    LE[*loopEntry].Tag = "L";
                    LE[*loopEntry].Tag += (char(loop_num + '0'));
                }
                else
                {
                    LE[*loopEntry].Tag = LE[UpperLoopEntry].Tag;
                    LE[*loopEntry].Tag += (char(loop_num + '0'));
                }
                get_Loop_Tag(*loopEntry, LE[*loopEntry].SuccLoop);
            }
        }

        //下面函数用于格式化打印循环

        //打印 indent 个空格, 缩进
        void print_indent(int indent)
        {
            for (int i = 0; i < indent; i++)
            {
                llvm::outs() << " ";
            }
        }

        void print_Loop_Function()
        {
            llvm::outs() << "{" << "\n";
            print_Loop(UpperMost_Loop);
            llvm::outs() << "}" << "\n";
        }

        void print_Loop(std::vector<DomTreeNode *> Entries)
        {
            for (auto loopEntry = Entries.begin(); loopEntry != Entries.end(); ++loopEntry)
            {
                int depth = (LE[*loopEntry].Tag.size() - 1);
                print_indent(4 * depth);
                llvm::outs() << "\"" << LE[*loopEntry].Tag << "\" : {" << "\n";
                print_indent(4 * (depth + 1));
                llvm::outs() << "\"depth\" : " << depth;
                if (!LE[*loopEntry].SuccLoop.empty())       //若下层循环不为空, 则打印 ',' 之后递归打印
                {
                    llvm::outs() << ",\n";
                    print_Loop(LE[*loopEntry].SuccLoop);
                }
                else                                        //否则仅换行
                {
                    llvm::outs() << "\n";
                }

                print_indent(4 * depth);
                if (loopEntry + 1 != Entries.end())         //若后续还有同层的循环, 则打印 ',' 并换行
                {
                    llvm::outs() << "},\n";
                }
                else                                        //否则直接换行
                {
                    llvm::outs() << "}\n";
                }
            }
        }

        bool runOnFunction(Function &F) override
        {
            if (skipFunction(F))
                return false;
            size_t BackEdgeNum = 0;
            DominatorTree DT(F);
            PostDominatorTree PDT(F);
            std::vector<DomTreeNode *> WorkList;
            LE.clear(); // map 初始化
            WorkList.push_back(DT.getRootNode());
            while (!WorkList.empty())
            {
                auto CurNode = WorkList.back();
                WorkList.pop_back();
                auto BB = CurNode->getBlock();
                for (auto sSucc = succ_begin(BB), eSucc = succ_end(BB); sSucc != eSucc; ++sSucc)
                {
                    auto SuccNode = DT.getNode(*sSucc);
                    if (DT.dominates(SuccNode, CurNode))
                    {
                        if (LE.find(SuccNode) == LE.end())
                        {
                            LoopEntry *entry = new LoopEntry;
                            entry->Level = SuccNode->getLevel();
                            entry->Pred.clear();
                            entry->Pred.push_back(CurNode);
                            entry->SuccLoop.clear();
                            entry->Tag.clear();
                            LE[SuccNode].UpperLoop = nullptr;
                            LE[SuccNode] = *entry;
                        }
                        else
                        {
                            LE[SuccNode].Pred.push_back(CurNode);
                        }
                        llvm::outs() << "Entry of a Loop: " << (SuccNode)->getBlock()->getName() << ", Level: " << SuccNode->getLevel() << "\n";
                        BackEdgeNum++;
                    }
                }
                WorkList.insert(WorkList.end(), CurNode->begin(), CurNode->end());
            }
            
            llvm::outs() << "Processing function " << F.getName() << ", number of Backedges is " << BackEdgeNum << "\n";
            std::error_code err;
            raw_fd_ostream outfile_dt(StringRef(F.getName().str() + "_dt.txt"), err);
            raw_fd_ostream outfile_pdt(StringRef(F.getName().str() + "_pdt.txt"), err);
            DT.print(outfile_dt);
            PDT.print(outfile_pdt);

            init_Map(DT);
            for (auto loopEntry = LE.begin(); loopEntry != LE.end(); ++loopEntry)
            {
                loopEntry->second.UpperLoop = get_Upper_Loop_Entry(DT, loopEntry->first);
            }
            
            get_Imm_Loop_Entry();
            get_Loop_Tag(nullptr, UpperMost_Loop);

#ifdef DEBUG
            //debug
            for (auto loopEntry = LE.begin(); loopEntry != LE.end(); ++loopEntry)
            {
                if (loopEntry->second.UpperLoop != nullptr)
                {
                    llvm::outs() << "Loop Entry " << loopEntry->first->getBlock()->getName() << "\'s Upper Loop Entry is: " << loopEntry->second.UpperLoop->getBlock()->getName() << "\n";
                }
            }

            for (auto loopEntry = LE.begin(); loopEntry != LE.end(); ++loopEntry)
            {
                //打印每个循环入口的直接下层循环
                llvm::outs() << "Loop Entry " << loopEntry->first->getBlock()->getName() << "\'s Succ Loop Entry contains: ";
                if (loopEntry->second.SuccLoop.empty())
                {
                    llvm::outs() << "NIL" << "\n";
                    continue;
                }
                for (auto SuccLoop = loopEntry->second.SuccLoop.begin(); SuccLoop != loopEntry->second.SuccLoop.end(); ++SuccLoop)
                {
                    llvm::outs() << (*SuccLoop)->getBlock()->getName() << " ";
                }
                llvm::outs() << "\n";
            }

            for (auto loopEntry = LE.begin(); loopEntry != LE.end(); ++loopEntry)
            {
                llvm::outs() << "Loop: " << loopEntry->first->getBlock()->getName() << "\'s Tag is :" << loopEntry->second.Tag << "\n";
            }
#endif
            print_Loop_Function();

            return true;
        }

        void getAnalysisUsage(AnalysisUsage &AU) const override
        {
            AU.addRequired<PostDominatorTreeWrapperPass>();
            AU.setPreservesCFG();
            AU.addPreserved<DominatorTreeWrapperPass>();
            AU.addPreserved<PostDominatorTreeWrapperPass>();
            AU.addPreserved<GlobalsAAWrapperPass>();
        }
    };

} // end anonymous namespace

char LSPass::ID = 0;
INITIALIZE_PASS(LSPass, "LSPass", "Loop search", false, false)

FunctionPass *llvm::createLSPass() { return new LSPass(); }